# -*- coding: UTF-8 -*-
import glob
import os
import cv2
import time
import imageio
import shutil
scene_name = "uv_lego"

base_dir = "/root/data/code/wxb/DNRSelect/logs/%s"%(scene_name)
exp_list = os.listdir(base_dir)
for exp_name in exp_list:
    # exp_name = "texdim_256_texfeature_16_camnum_16_first_down1_imgresize800_raytraced_lr0.001_b1_e1000_dropcam0.0_trainratio_0.1_2023-05-08_17-33-15"
    log_dir = os.path.join(base_dir, exp_name)
    log_dir_visual = os.path.join(log_dir, "visual")

    # 图片合成视频
    def picvideo(path, size, output):
        filelist = os.listdir(path)  # 获取该目录下的所有文件名
        filelist.sort()  ##文件名按数字排序

        fps = 30
        file_path = os.path.join(output, "video.avi")  # 导出路径
        fourcc = cv2.VideoWriter_fourcc('I','4','2','0')  # 不同视频编码对应不同视频格式（例：'I','4','2','0' 对应avi格式）

        video = cv2.VideoWriter(file_path, fourcc, fps, size)

        for item in filelist:
            if "test" in item:
                # print(item)
                img = cv2.imread(os.path.join(path,item))  # 使用opencv读取图像，直接返回numpy.ndarray 对象，通道顺序为BGR ，注意是BGR，通道值默认范围0-255。
                video.write(img)  # 把图片写进视频

        video.release()  # 释放

    def picgif(path, size, output):
        filelist = os.listdir(path)  # 获取该目录下的所有文件名
        filelist.sort()  ##文件名按数字排序
        fps = 30
        file_path = os.path.join(output, "video.gif")  # 导出路径

        with imageio.get_writer(uri=file_path, mode='I', fps=fps) as writer:
            for item in filelist:
                if "test" in item:
                    img = imageio.imread(os.path.join(path,item))
                    # img = img[:, 1600:2400, :]
                    # img = img[:, 2400:, :]
                    # img = cv2.resize(img, size)
                    writer.append_data(img)
    print(log_dir_visual)
    picvideo(log_dir_visual, (3200,800), log_dir)
    picgif(log_dir_visual, (400,400), log_dir)
    os.remove(os.path.join(log_dir, "model.pth"))
    shutil.rmtree(log_dir_visual)
    shutil.rmtree(os.path.join(log_dir, "scripts"))

